module cubed_sphere_math_mod
  use kinds_mod
  use mesh_const_mod
  use mesh_math_mod

  implicit none

  private

  public locate_domain_interface
  public cube_to_sphere_interface
  public sphere_to_cube_interface
  public horizontal_metrics_interface
  public vertical_metrics_interface
  public equiangular_locate_domain
  public equiangular_cube_to_sphere
  public equiangular_sphere_to_cube
  public equiangular_get_jab
  public equiangular_metrics
  public vertical_metrics
  public cov_to_spherev
  public contrav_to_spherev
  public spherev_to_contrav
  public contrav_to_cov
  public cov_to_contrav

  interface
    subroutine locate_domain_interface(lon, lat, idom1, x1, y1, idom2, x2, y2)
      import r8
      real(r8), intent(in ) :: lon
      real(r8), intent(in ) :: lat
      integer , intent(out) :: idom1
      real(r8), intent(out) :: x1
      real(r8), intent(out) :: y1
      integer , intent(out), optional :: idom2
      real(r8), intent(out), optional :: x2
      real(r8), intent(out), optional :: y2
    end subroutine locate_domain_interface
    subroutine cube_to_sphere_interface(idom, x, y, lon, lat)
      import r8
      integer , intent(in ) :: idom
      real(r8), intent(in ) :: x
      real(r8), intent(in ) :: y
      real(r8), intent(out) :: lon
      real(r8), intent(out) :: lat
    end subroutine cube_to_sphere_interface
    subroutine sphere_to_cube_interface(idom, lon, lat, x, y)
      import r8
      integer , intent(in ) :: idom
      real(r8), intent(in ) :: lon
      real(r8), intent(in ) :: lat
      real(r8), intent(out) :: x
      real(r8), intent(out) :: y
    end subroutine sphere_to_cube_interface
    subroutine horizontal_metrics_interface( idom, r, x, y, lon, lat, G, iG, A, iA, J, CS )
      import r16
      integer , intent(in ) :: idom
      real(r16), intent(in ) :: r
      real(r16), intent(in ) :: x
      real(r16), intent(in ) :: y
      real(r16), intent(in ) :: lon
      real(r16), intent(in ) :: lat
      real(r16), intent(out) :: G (3,3)
      real(r16), intent(out) :: iG(3,3)
      real(r16), intent(out) :: A (3,3)
      real(r16), intent(out) :: iA(3,3)
      real(r16), intent(out) :: J
      real(r16), intent(out) :: CS(3,3,3)
      end subroutine horizontal_metrics_interface
    subroutine vertical_metrics_interface( Ra, Rb, Rr, G, iG, A, iA, J )
      import r16
      real(r16), intent(in ) :: Ra
      real(r16), intent(in ) :: Rb
      real(r16), intent(in ) :: Rr
      real(r16), intent(out) :: G (3,3)
      real(r16), intent(out) :: iG(3,3)
      real(r16), intent(out) :: A (3,3)
      real(r16), intent(out) :: iA(3,3)
      real(r16), intent(out) :: J
      end subroutine vertical_metrics_interface
  end interface

contains

  subroutine equiangular_locate_domain(lon, lat, idom1, x1, y1, idom2, x2, y2)

    real(r8), intent(in ) :: lon
    real(r8), intent(in ) :: lat
    integer , intent(out) :: idom1
    real(r8), intent(out) :: x1
    real(r8), intent(out) :: y1
    integer , intent(out), optional :: idom2
    real(r8), intent(out), optional :: x2
    real(r8), intent(out), optional :: y2

    real(r8), parameter :: eps = 1.0e-14_r8

    if (lon >= 7 * pi0p25 .or.  lon <=     pi0p25) then
      idom1 = 1
      if (present(idom2)) idom2 = merge(2, merge(4, idom1, abs(lon - 7 * pi0p25) <= eps), abs(lon -     pi0p25) <= eps)
    end if
    if (lon >=     pi0p25 .and. lon <= 3 * pi0p25) then
      idom1 = 2
      if (present(idom2)) idom2 = merge(3, merge(1, idom1, abs(lon -     pi0p25) <= eps), abs(lon - 3 * pi0p25) <= eps)
    end if
    if (lon >= 3 * pi0p25 .and. lon <= 5 * pi0p25) then
      idom1 = 3
      if (present(idom2)) idom2 = merge(4, merge(2, idom1, abs(lon - 3 * pi0p25) <= eps), abs(lon - 5 * pi0p25) <= eps)
    end if
    if (lon >= 5 * pi0p25 .and. lon <= 7 * pi0p25) then
      idom1 = 4
      if (present(idom2)) idom2 = merge(1, merge(3, idom1, abs(lon - 5 * pi0p25) <= eps), abs(lon - 7 * pi0p25) <= eps)
    end if

    call equiangular_sphere_to_cube(idom1, lon, lat, x1, y1)

    if (abs(y1 + pi0p25) <= eps) then
      if (present(idom2)) idom2 = 6
    else if (abs(y1 - pi0p25) <= eps) then
      if (present(idom2)) idom2 = 5
    else if (y1 <= -pi0p25) then
      idom1 = 6
      if (present(idom2))idom2 = 6
    else if (y1 >= pi0p25) then
      idom1 = 5
      if (present(idom2))idom2 = 5
    end if

    call equiangular_sphere_to_cube(idom1, lon, lat, x1, y1)
    if (present(idom2)) call equiangular_sphere_to_cube(idom2, lon, lat, x2, y2)

  end subroutine equiangular_locate_domain

  subroutine equiangular_cube_to_sphere(idom, x, y, lon, lat)

    integer , intent(in ) :: idom
    real(r8), intent(in ) :: x
    real(r8), intent(in ) :: y
    real(r8), intent(out) :: lon
    real(r8), intent(out) :: lat

    real(r8) a, b

    select case (idom)
    case (1:4)
      lon = x + (idom - 1) * pi0p5
      lat = atan2(tan(y) * cos(x), 1.0_r8)
    case (5)
      a = tan(x)
      b = tan(y)
      lon = atan2(a, -b)
      lat = atan2(1.0_r8, sqrt(a**2 + b**2))
    case (6)
      a = tan(x)
      b = tan(y)
      lon = atan2(a, b)
      lat = -atan2(1.0_r8, sqrt(a**2 + b**2))
    end select

    if (lon < 0.0) lon = lon + pi2
    if (lon > pi2) lon = lon - pi2

  end subroutine equiangular_cube_to_sphere

  subroutine equiangular_sphere_to_cube(idom, lon, lat, x, y)

    integer , intent(in ) :: idom
    real(r8), intent(in ) :: lon
    real(r8), intent(in ) :: lat
    real(r8), intent(out) :: x
    real(r8), intent(out) :: y

    select case (idom)
    case (1:4)
      x = atan(tan(lon - (idom - 1) * pi0p5))
      y = atan(tan(lat) / cos(lon - (idom - 1) * pi0p5))
    case (5)
      x = atan( sin(lon) / tan(lat))
      y = atan(-cos(lon) / tan(lat))
    case (6)
      x = atan(-sin(lon) / tan(lat))
      y = atan(-cos(lon) / tan(lat))
    end select

    ! Avoid x and y outside [-pi/4,pi/4].
    if (abs(abs(x) - pi0p25) < 1.0e-15_r8) x = sign(1.0_r8, x) * pi0p25
    if (abs(abs(y) - pi0p25) < 1.0e-15_r8) y = sign(1.0_r8, y) * pi0p25

  end subroutine equiangular_sphere_to_cube

  ! subroutine equiangular_metrics( idom, r, x, y, lon, lat, G, iG, A, iA, J, CS )
  !   integer , intent(in ) :: idom
  !   real(r16), intent(in ) :: r
  !   real(r16), intent(in ) :: x
  !   real(r16), intent(in ) :: y
  !   real(r16), intent(in ) :: lon
  !   real(r16), intent(in ) :: lat
  !   real(r16), intent(out) :: G(3,3), iG(3,3)
  !   real(r16), intent(out) :: A(3,3), iA(3,3)
  !   real(r16), intent(out) :: J
  !   real(r16), intent(out) :: CS(3,3,3)

  !   real(r16) sinx, siny
  !   real(r16) cscx, cscy
  !   real(r16) cosx, cosy
  !   real(r16) secx, secy
  !   real(r16) tanx, tany
  !   real(r16) cotx, coty

  !   integer i
  !   integer ierr

  !   sinx = sin(x)
  !   siny = sin(y)
  !   cscx = 1. / sinx
  !   cscy = 1. / siny
  !   cosx = cos(x)
  !   cosy = cos(y)
  !   secx = 1. / cosx
  !   secy = 1. / cosy
  !   tanx = tan(x)
  !   tany = tan(y)
  !   cotx = 1. / tanx
  !   coty = 1. / tany

  !   A = 0
  !   select case (idom)
  !   case (1:4)
  !     A(1,1) = r / sqrt( 1. + cosx**2 * tany**2 )
  !     A(1,2) = 0
  !     A(2,1) = - r * coty * sinx / ( cosx**2 + coty**2 )
  !     A(2,2) = r * cosx * secy**2 / ( 1 + cosx**2 * tany**2 )
  !   case (5)
  !     A(1,1) =-r * coty * secx**2 / ( ( 1 + coty**2 * tanx**2 ) * sqrt( 1 + 1 / ( tanx**2 + tany**2 ) ) )
  !     A(1,2) = r * cscy**2 * tanx / ( ( 1 + coty**2 * tanx**2 ) * sqrt( 1 + 1 / ( tanx**2 + tany**2 ) ) )
  !     A(2,1) =-r * secx**2 * tanx / ( ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) )
  !     A(2,2) =-r * secy**2 * tany / ( ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) )
  !   case (6)
  !     A(1,1) = r * coty * secx**2 / ( ( 1 + coty**2 * tanx**2 ) * sqrt( 1 + 1 / ( tanx**2 + tany**2 ) ) )
  !     A(1,2) =-r * cscy**2 * tanx / ( ( 1 + coty**2 * tanx**2 ) * sqrt( 1 + 1 / ( tanx**2 + tany**2 ) ) )
  !     A(2,1) = r * secx**2 * tanx / ( ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) )
  !     A(2,2) = r * secy**2 * tany / ( ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) )
  !   end select
  !   A(3,3) = 1

  !   ! iA = 0
  !   ! select case (idom)
  !   ! case (1:4)
  !   !   iA(1,1) = sqrt( ( 1 + cosx**2 * tany ) ) / r
  !   !   iA(1,2) = 0
  !   !   iA(2,1) = cosy * siny * tanx * iA(1,1)
  !   !   iA(2,2) = cosx * cosy**2 * ( secx**2 + tany**2 ) / r
  !   ! case (5)
  !   !   iA(1,1) =-cosx**2 * tany * sqrt( 1. + 1. / ( tanx**2 + tany**2 ) ) / r
  !   !   iA(1,2) =-cosx * cscy**2 * sinx * ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) / ( r * ( secy**2 + cscy**2 * tanx**2 ) )
  !   !   iA(2,1) =-iA(1,1)
  !   !   iA(2,2) =-coty * ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) / ( r * ( secy**2 + cscy**2 * tanx**2 ) )
  !   ! case (6)
  !   !   iA(1,1) = cosx**2 * tany * sqrt( 1. + 1. / ( tanx**2 + tany**2 ) ) / r
  !   !   iA(1,2) = cosx * cscy**2 * sinx * ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) / ( r * ( secy**2 + cscy**2 * tanx**2 ) )
  !   !   iA(2,1) =-iA(1,1)
  !   !   iA(2,2) = coty * ( secy**2 + tanx**2 ) * sqrt( tanx**2 + tany**2 ) / ( r * ( secy**2 + cscy**2 * tanx**2 ) )
  !   ! end select
  !   ! iA(3,3) = 1
    
  !   call inverse_matrix(A, iA, ierr)
  !   if(ierr/=0)then
  !     print*,idom,x,y,lon,lat,r
  !     do i = 1,3
  !       print*,A(i,:)
  !     enddo
  !     stop 'Calculate inverse A failed in equiangular_metrics'
  !   endif
    
  !   G = matmul(transpose(A),A)
  !   call inverse_matrix(G, iG, ierr)
  !   if(ierr/=0)then
  !     print*,idom,x,y,lon,lat,r
  !     do i = 1,3
  !       print*,A(i,:)
  !     enddo
  !     stop 'Calculate inverse G failed in equiangular_metrics'
  !   endif
    
  !   J = sqrt( det(G) )
    
  !   CS = equiangular_Christoffel_symbol(x,y,r)
    
  ! end subroutine equiangular_metrics

  subroutine equiangular_get_jab(jab,x,y,r)
    real(r_kind), intent(in   ) :: x,y,r
    real(r_kind), intent(  out) :: jab
    real :: rho
    rho = sqrt( 1._r_kind + tan(x)**2 + tan(y)**2 )
    jab = r**2 / ( cos(x)**2 * cos(y)**2 * rho**3 )
  end subroutine equiangular_get_jab

  subroutine equiangular_metrics( idom, r, x, y, lon, lat, G, iG, A, iA, J, CS )
    integer , intent(in ) :: idom
    real(r16), intent(in ) :: r
    real(r16), intent(in ) :: x
    real(r16), intent(in ) :: y
    real(r16), intent(in ) :: lon
    real(r16), intent(in ) :: lat
    real(r16), intent(out) :: G(3,3), iG(3,3)
    real(r16), intent(out) :: A(3,3), iA(3,3)
    real(r16), intent(out) :: J
    real(r16), intent(out) :: CS(3,3,3)

    real(r16) rho, tanx, tany, cosx2, cosy2
    real(r16) lon0, lat0, e, b, c, d
    integer i
    integer ierr

    A = 0
    select case (idom)
    case (1:4)
      lon0 = lon - (idom - 1) * pi / 2
      lat0 = lat
      e = sin(lon0); b = cos(lon0); c = sin(lat0); d = cos(lat0)
      A(1,1) = d
      A(1,2) = 0
      A(2,1) = -c * d * e / b
      A(2,2) = b * d**2 + c**2 / b
    case (5)
      lon0 = lon
      lat0 = lat
      e = sin(lon0); b = cos(lon0); c = sin(lat0); d = cos(lat0)
      rho = 1 + e**2 * d**2 / c**2
      A(1,1) = b * c * rho
      A(2,1) = -c**2 * e * rho
      rho = 1 + b**2 * d**2 / c**2
      A(1,2) = e * c * rho
      A(2,2) = b * c**2 * rho
    case (6)
      lon0 = lon
      lat0 = lat
      e = sin(lon0); b = cos(lon0); c = sin(lat0); d = cos(lat0)
      rho = 1 + e**2 * d**2 / c**2
      A(1,1) = -b * c * rho
      A(2,1) = c**2 * e * rho
      rho = 1 + b**2 * d**2 / c**2
      A(1,2) = e * c * rho
      A(2,2) = b * c**2 * rho
    end select

    A = A * r
    A(3,1) = 0; A(3,2) = 0; A(3,3) = 1
    
    call inverse_matrix(A, iA, ierr)
    if(ierr/=0)then
      print*,idom,x,y,lon,lat,r
      do i = 1,3
        print*,A(i,:)
      enddo
      stop 'Calculate inverse A failed in equiangular_metrics'
    endif
    
    G = matmul(transpose(A),A)
    call inverse_matrix(G, iG, ierr)
    if(ierr/=0)then
      print*,idom,x,y,lon,lat,r
      do i = 1,3
        print*,A(i,:)
      enddo
      stop 'Calculate inverse G failed in equiangular_metrics'
    endif
    
    J = sqrt( det(G) )
    
    CS = equiangular_Christoffel_symbol(x,y,r)
    
  end subroutine equiangular_metrics
  
  subroutine vertical_metrics( Ra, Rb, Rr, G, iG, A, iA, J )
    real(r16), intent(in ) :: Ra
    real(r16), intent(in ) :: Rb
    real(r16), intent(in ) :: Rr
    real(r16), intent(out) :: G (3,3)
    real(r16), intent(out) :: iG(3,3)
    real(r16), intent(out) :: A (3,3)
    real(r16), intent(out) :: iA(3,3)
    real(r16), intent(out) :: J

    integer i
    integer ierr

    ! A(1,1) =  1; A(1,2) =  0; A(1,3) = 0
    ! A(2,1) =  0; A(2,2) =  1; A(2,3) = 0
    ! A(3,1) = Ra; A(3,2) = Rb; A(3,3) = Rr
    
    ! call inverse_matrix(A, iA, ierr)
    ! if(ierr/=0)then
    !   print*,Ra,Rb,Rr
    !   do i = 1,3
    !     print*,A(i,:)
    !   enddo
    !   stop 'Calculate inverse A failed in vertical_metrics'
    ! endif
    
    ! G = matmul(transpose(A),A)
    ! call inverse_matrix(G, iG, ierr)
    ! if(ierr/=0)then
    !   print*,Ra,Rb,Rr
    !   do i = 1,3
    !     print*,A(i,:)
    !   enddo
    !   stop 'Calculate inverse G failed in vertical_metrics'
    ! endif
    
    ! J = sqrt( det(G) )

    iG(1,3) = -Ra / Rr
    iG(2,3) = -Rb / Rr
    J = Rr
    
  end subroutine vertical_metrics
  
  subroutine cov_to_spherev(A, iG, u1, u2, v1, v2)

    real(r8), intent(in ) :: A (2,2)
    real(r8), intent(in ) :: iG(2,2)
    real(r8), intent(in ) :: u1 ! Covariant vector component
    real(r8), intent(in ) :: u2 ! Covariant vector component
    real(r8), intent(out) :: v1 ! Spherical vector component
    real(r8), intent(out) :: v2 ! Spherical vector component
    
    real(r8) A_iG(2,2)
    
    A_iG = matmul(A, iG)
    
    v1 = A_iG(1,1) * u1 + A_iG(1,2) * u2
    v2 = A_iG(2,1) * u1 + A_iG(2,2) * u2

  end subroutine cov_to_spherev
  
  subroutine spherev_to_contrav(iA, u1, u2, v1, v2)

    real(r8), intent(in ) :: iA(2,2)
    real(r8), intent(in ) :: u1 ! Spherical vector component
    real(r8), intent(in ) :: u2 ! Spherical vector component
    real(r8), intent(out) :: v1 ! Contravariant vector component
    real(r8), intent(out) :: v2 ! Contravariant vector component
          
    v1 = iA(1,1) * u1 + iA(1,2) * u2
    v2 = iA(2,1) * u1 + iA(2,2) * u2

  end subroutine spherev_to_contrav

  subroutine contrav_to_spherev(A, u1, u2, v1, v2)

    real(r8), intent(in ) :: A(2,2)
    real(r8), intent(in ) :: u1 ! Contravariant vector component
    real(r8), intent(in ) :: u2 ! Contravariant vector component
    real(r8), intent(out) :: v1 ! Spherical vector component
    real(r8), intent(out) :: v2 ! Spherical vector component
    
    v1 = A(1,1) * u1 + A(1,2) * u2
    v2 = A(2,1) * u1 + A(2,2) * u2

  end subroutine contrav_to_spherev

  subroutine contrav_to_cov(G, u1, u2, v1, v2)

    real(r8), intent(in ) :: G(2,2)
    real(r8), intent(in ) :: u1 ! Contravariant vector component
    real(r8), intent(in ) :: u2 ! Contravariant vector component
    real(r8), intent(out) :: v1 ! Covariant vector component
    real(r8), intent(out) :: v2 ! Covariant vector component
    
    v1 = G(1,1) * u1 + G(1,2) * u2
    v2 = G(2,1) * u1 + G(2,2) * u2

  end subroutine contrav_to_cov
  
  subroutine cov_to_contrav(iG, u1, u2, v1, v2)

    real(r8), intent(in ) :: iG(2,2)
    real(r8), intent(in ) :: u1 ! Covariant vector component
    real(r8), intent(in ) :: u2 ! Covariant vector component
    real(r8), intent(out) :: v1 ! Contravariant vector component
    real(r8), intent(out) :: v2 ! Contravariant vector component
    
    v1 = iG(1,1) * u1 + iG(1,2) * u2
    v2 = iG(2,1) * u1 + iG(2,2) * u2
    
  end subroutine cov_to_contrav
  
  function equiangular_Christoffel_symbol(alpha,beta,r) result(CS)
    real(r16) :: CS(3,3,3)
    real(r16), intent(in) :: alpha
    real(r16), intent(in) :: beta
    real(r16), intent(in) :: r
    
    ! (Gamma_jk)^i
    !  j,k,i
    CS(1, 1, 1) = -(2*sin(2*alpha)*(cos(2*beta) - 1))/(cos(2*alpha) + cos(2*beta) - cos(2*alpha)*cos(2*beta) + 3)           
    CS(1, 2, 1) = -(sin(beta)*(cos(2*alpha) + 1))/(cos(beta)*(cos(2*alpha) + cos(beta)**2 - cos(2*alpha)*cos(beta)**2 + 1))   
    CS(1, 3, 1) = 1/R                                                                                                       
    CS(2, 1, 1) = -(sin(beta)*(cos(2*alpha) + 1))/(cos(beta)*(cos(2*alpha) + cos(beta)**2 - cos(2*alpha)*cos(beta)**2 + 1))   
    CS(2, 2, 1) = 0                                                                                                         
    CS(2, 3, 1) = 0                                                                                                         
    CS(3, 1, 1) = 1/R                                                                                                       
    CS(3, 2, 1) = 0                                                                                                         
    CS(3, 3, 1) = 0                                                                                                         
    CS(1, 1, 2) = 0                                                                                                         
    CS(1, 2, 2) = -(sin(alpha)*(cos(2*beta) + 1))/(cos(alpha)*(cos(2*beta) + cos(alpha)**2 - cos(2*beta)*cos(alpha)**2 + 1))  
    CS(1, 3, 2) = 0                                                                                                         
    CS(2, 1, 2) = -(sin(alpha)*(cos(2*beta) + 1))/(cos(alpha)*(cos(2*beta) + cos(alpha)**2 - cos(2*beta)*cos(alpha)**2 + 1))  
    CS(2, 2, 2) = -(2*sin(2*beta)*(cos(2*alpha) - 1))/(cos(2*alpha) + cos(2*beta) - cos(2*alpha)*cos(2*beta) + 3)           
    CS(2, 3, 2) = 1/R                                                                                                       
    CS(3, 1, 2) = 0                                                                                                         
    CS(3, 2, 2) = 1/R                                                                                                       
    CS(3, 3, 2) = 0                                                                                                         
    CS(1, 1, 3) = -(R*(tan(alpha)**2 + 1)**2*(tan(beta)**2 + 1))/(tan(alpha)**2 + tan(beta)**2 + 1)**2                            
    CS(1, 2, 3) = (R*tan(alpha)*tan(beta)*(tan(alpha)**2 + 1)*(tan(beta)**2 + 1))/(tan(alpha)**2 + tan(beta)**2 + 1)**2          
    CS(1, 3, 3) = 0                                                                                                         
    CS(2, 1, 3) = (R*tan(alpha)*tan(beta)*(tan(alpha)**2 + 1)*(tan(beta)**2 + 1))/(tan(alpha)**2 + tan(beta)**2 + 1)**2          
    CS(2, 2, 3) = -(R*(tan(alpha)**2 + 1)*(tan(beta)**2 + 1)**2)/(tan(alpha)**2 + tan(beta)**2 + 1)**2                            
    CS(2, 3, 3) = 0                                                                                                         
    CS(3, 1, 3) = 0                                                                                                         
    CS(3, 2, 3) = 0                                                                                                         
    CS(3, 3, 3) = 0                                                                                                         
  end function equiangular_Christoffel_symbol
  
end module cubed_sphere_math_mod
